import os.path
import numpy as np
import torch
from copy import deepcopy
from tensorboardX import SummaryWriter
from tqdm import tqdm
import torch.nn as nn
import random
from medpy import metric
from scipy.ndimage import zoom
import torch.nn as nn
import os
from tqdm import tqdm
from torchvision.utils import make_grid
import torch.nn.functional as F
from utils import loadyaml, _get_logger, mk_path, get_current_consistency_weight, DiceLoss, update_ema_variables
from utils import build_lr_scheduler, build_optimizer, Med_Sup_Loss, Dense_Loss, BoxMaskGenerator
from model import build_model
from datasets import build_loader


def main():

    path = r"config/ddpm_800k_cifar.yaml"
    root = os.path.dirname(os.path.realpath(__file__))  # 获取绝对路径
    args = loadyaml(os.path.join(root, path))  # 加载yaml
    if args.cuda:
        args.device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
    else:
        args.device = torch.device("cpu")

    root = os.path.dirname(os.path.realpath(__file__))  # 获取绝对路径
    args.save_path = os.path.join(root, args.save_path)
    mk_path(args.save_path)  # 创建文件保存位置
    mk_path(os.path.join(args.save_path, "tensorboardX")) # 创建 tensorboardX日志保存位置
    mk_path(os.path.join(args.save_path, "model"))  # 创建模型保存位置

    args.model_save_path = os.path.join(args.save_path, "model", "model.pth")
    args.writer = SummaryWriter(os.path.join(args.save_path, "tensorboardX"))
    args.logger = _get_logger(os.path.join(args.save_path, "log.log"), "info")
    
    args.tqdm = os.path.join(args.save_path, "tqdm.log")
    torch.manual_seed(args.seed)  # 设置随机种子
    torch.cuda.manual_seed(args.seed)
    np.random.seed(args.seed)

    torch.backends.cudnn.deterministic = False  # 单卡的不需要分布式
    torch.backends.cudnn.benchmark = True  # 寻找最佳 的训练路径

    train_loader,test_loader = build_loader(args)  # 构建数据集
    args.epochs = args.total_itrs // len(train_loader)  # 设置模型epoch
    args.logger.info("==========> train_loader length:{}".format(len(train_loader.dataset)))
    args.logger.info("==========> test_loader length:{}".format(len(test_loader.dataset)))
    args.logger.info("==========> epochs length:{}".format(args.epochs))

    # step 1: 构建模型
    model = build_model(args=args).to(device=args.device)  # 创建模型

    # step 2: 训练模型
    DDPM(model, train_loader, test_loader, args)


def DDPM(model,  train_loader, test_loader, args):

    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
    cur_itrs=0
    pbar = tqdm(total=args.total_itrs)

    for epoch in range(args.epochs):
        train_loss=0.0
        model.train()

        for x, y in train_loader:
            cur_itrs+=1

            x = x.to(args.device)
            y = y.to(args.device)


            if args.use_labels:
                loss = model(x, y)
            else:
                loss = model(x)
            
            train_loss += loss.item()

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            args.writer.add_scalar('ddpm/loss', loss.item(), cur_itrs)

            model.update_ema()

            if cur_itrs % args.step_size == 0:
                ddpm_eval(cur_itrs=cur_itrs,model=model,test_loader=test_loader,args=args)
                model.train()

            if cur_itrs > args.total_itrs:
                return
            
            pbar.update(1)

        args.logger.info("Train  [{}/{} ({:.0f}%)]\t loss: {:.5f}".format(cur_itrs, args.total_itrs,
                             100. * cur_itrs / args.total_itrs, train_loss / len(train_loader)))


def ddpm_eval(cur_itrs,model,test_loader,args):

    test_loss = 0.0

    with torch.no_grad():
        model.eval()
        for x, y in tqdm(test_loader):
            x = x.to(args.device)
            y = y.to(args.device)

            if args.use_labels:
                loss = model(x, y)
            else:
                loss = model(x)

            test_loss += loss.item()
    
    if args.use_labels:
        samples = model.sample(64, args.device, y=torch.arange(64, device=args.device))
    else:
        samples = model.sample(64, args.device)
    
    # samples = ((samples + 1) / 2).clip(0, 1)

    img=make_grid(samples,nrow=8,normalize=True,scale_each=True)
    args.writer.add_image('ddpm_test/Image', img, cur_itrs)
    args.writer.add_scalar('ddpm_test/test_loss', test_loss, cur_itrs)




if __name__ == '__main__':
    main()
